#!/usr/bin/env python3
"""
IGRAG Main Entry Point
图像检索增强生成主程序

功能：整合所有模块，实现端到端的图像描述生成
"""

import os
import sys
import yaml
import logging
import argparse
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

from core.retriever import ImageRetriever
from core.generator import CaptionGenerator
from utils.image_utils import load_image


def setup_logging(log_level: str = "ERROR"):
    """
    设置日志配置
    
    Args:
        log_level: 日志级别
    """
    logging.basicConfig(
        level=getattr(logging, log_level.upper()),
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )


def load_config(config_path: str) -> dict:
    """
    加载配置文件
    
    Args:
        config_path: 配置文件路径
        
    Returns:
        dict: 配置字典
    """
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        return config
    except Exception as e:
        logging.error(f"加载配置文件失败: {e}")
        raise


def process_image(image_path: str, config: dict) -> str:
    """
    处理单张图像，生成描述
    
    Args:
        image_path: 图像路径
        config: 配置字典
        
    Returns:
        str: 生成的图像描述
    """
    try:
        # 加载图像
        image = load_image(image_path)
        if image is None:
            raise ValueError(f"无法加载图像: {image_path}")
        
        # 初始化检索器和生成器
        retriever = ImageRetriever(config)
        generator = CaptionGenerator(config)
        
        # 检索相似图像并获取描述
        retrieved_captions = retriever.get_retrieved_captions(image)
        
        # 构建提示词
        prompt = generator.build_prompt(retrieved_captions)
        
        # 生成最终描述
        generated_caption = generator.generate_caption(prompt)
        
        return generated_caption
        
    except Exception as e:
        logging.error(f"处理图像失败: {e}")
        raise


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="IGRAG - 图像检索增强生成")
    parser.add_argument(
        "--input", 
        type=str, 
        required=True,
        help="输入图像路径"
    )
    parser.add_argument(
        "--config", 
        type=str, 
        default="configs/config.yaml",
        help="配置文件路径"
    )
    parser.add_argument(
        "--log-level", 
        type=str, 
        default="ERROR",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        help="日志级别"
    )
    
    args = parser.parse_args()
    
    try:
        # 设置日志
        setup_logging(args.log_level)
        logger = logging.getLogger(__name__)
        
        # 检查输入文件
        if not os.path.exists(args.input):
            logger.error(f"输入图像文件不存在: {args.input}")
            sys.exit(1)
        
        # 加载配置
        config_path = project_root / args.config
        if not config_path.exists():
            logger.error(f"配置文件不存在: {config_path}")
            sys.exit(1)
        
        config = load_config(str(config_path))
        
        # 处理图像
        logger.info(f"开始处理图像: {args.input}")
        generated_caption = process_image(args.input, config)
        
        # 输出结果
        print("=" * 50)
        print("IGRAG 图像描述生成结果")
        print("=" * 50)
        print(f"输入图像: {args.input}")
        print(f"生成描述: {generated_caption}")
        print("=" * 50)
        
    except Exception as e:
        logging.error(f"程序执行失败: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()
